import csv
import pickle
import sys
import os
from lib import tools as my_tools

from tqdm import tqdm

# def dataset2csv(index:str):
#     """
#     将指定序号的数据集转换为CSV格式, 并保存到文件中。
#     python dataset2csv.py index
#     index: 0-based
#     """
#     # 加载数据
#     data_set = load_data()  

#     #将字符串index转换为对应的整数
#     try:
#         index_value = int(index)
#     except ValueError:
#         print(f"Invalid index: {index}, setting to default value 0")
#         index_value = 0


#     if index_value >= len(data_set):
#         index_value = len(data_set) - 1
#         index = str(index_value)
#         print(f"Index out of range, set to last index value: {index_value}") 

#     print(f"Waiting for sorting dataset")
#     with open('dataset'+ index +'.csv', 'w', newline='', encoding='utf-8') as csvfile:
#         writer = csv.writer(csvfile)
#         # 写入表头
#         writer.writerow(['D_value', 'D1', 'D2', 'D3', 'v_zd', 'omega_x', 'omega_y', 'omega_z', 'p_x', 'p_y', 'c_z', 'ev_z', 'gv_z', 'n3'])
       
#         if len(data_set[index_value]) > 0:
#             try:
#                 data_set[index_value].sort(key=lambda x: x[0])
#             except Exception as e:
#                 print(f"Error while sorting data_set[{index_value}]: {e}")
#                 print(f"data_set[{index_value}] is {data_set[index_value]}")
#             writer.writerows(data_set[index_value])
#     print(f"Dataset sorted and saved to dataset{index}.csv")

def dataset2csv():
    # 加载数据
    data_set = my_tools.load_dataset('found_data_set.pickle')
    print(f"dataset length : {len(data_set)}")

    # 检查目标文件夹是否存在
    folder_name = 'subdataset'
    if not os.path.exists(folder_name):
        os.makedirs(folder_name)
    
    for i in tqdm(range(len(data_set))):
        with open(folder_name+'/dataset' + str(i) + '.csv', 'w', newline='', encoding='utf-8') as csvfile:
            writer = csv.writer(csvfile)
            # 写入表头
            writer.writerow(['D_value', 'D1', 'D2', 'D3', 'v_zd', 'omega_x', 'omega_y', 'omega_z', 'p_x', 'p_y', 'c_z', 'ev_z', 'gv_z', 'n3'])
        
            if len(data_set[i]) > 0:
                try:
                    data_set[i].sort(key=lambda x: x[0])
                except Exception as e:
                    print(f"Error while sorting data_set[{i}]: {e}")
                    print(f"data_set[{i}] is {data_set[i]}")
                writer.writerows(data_set[i])
                tqdm.write(f"Dataset with item len {len(data_set[i])} sorted and saved to dataset{i}.csv")
            else:
                tqdm.write(f"Empty Dataset skipped: dataset{i}.csv")

def dataset_attitude2csv():
    # 检查目标文件夹是否存在
    folder_name = 'subdataset_attitude'
    if not os.path.exists(folder_name):
        os.makedirs(folder_name)


    dataset_reader = my_tools.load_dataset_chunk('found_data_set_attitude.pickle')
    j = 0
    for data_set in tqdm(dataset_reader, desc="Processing data chunks"):
        for i in range(len(data_set)):
            # tqdm.write(f"dataset length : {len(data_set[i])}")
            j+=1
            with open(folder_name+'/dataset' + str(j) + '.csv', 'w', newline='', encoding='utf-8') as csvfile:
                writer = csv.writer(csvfile)
                # 写入表头
                writer.writerow(['D_value', 'omega_x', 'omega_y', 'omega_z', 'n3_x', 'n3_y', 'n3_z', 'n3d_x', 'n3d_y', 'n3d_z'])

                if len(data_set[i]) > 0:
                    try:
                        data_set[i].sort(key=lambda x: x[0])
                    except Exception as e:
                        tqdm.write(f"Error while sorting data_set[{i}]: {e}")
                        tqdm.write(f"data_set[{i}] is {data_set[i]}")
                    writer.writerows(data_set[i])
                    # tqdm.write(f"Dataset with item len {len(data_set[i])} sorted and saved to dataset{i}.csv")
                else:
                    tqdm.write(f"Empty Dataset skipped: dataset{i}.csv")

    # # 加载数据
    # data_set = my_tools.load_dataset('found_data_set_attitude.pickle')
    # print(f"dataset length : {len(data_set)}")
    
    # for i in tqdm(range(len(data_set))):
    #     with open(folder_name+'/dataset' + str(i) + '.csv', 'w', newline='', encoding='utf-8') as csvfile:
    #         writer = csv.writer(csvfile)
    #         # 写入表头
    #         writer.writerow(['D_value', 'omega_x', 'omega_y', 'omega_z', 'n3', 'n3d'])
        
    #         if len(data_set[i]) > 0:
    #             try:
    #                 data_set[i].sort(key=lambda x: x[0])
    #             except Exception as e:
    #                 print(f"Error while sorting data_set[{i}]: {e}")
    #                 print(f"data_set[{i}] is {data_set[i]}")
    #             writer.writerows(data_set[i])
    #             tqdm.write(f"Dataset with item len {len(data_set[i])} sorted and saved to dataset{i}.csv")
    #         else:
    #             tqdm.write(f"Empty Dataset skipped: dataset{i}.csv")

# def load_data():
#     print(f"Waiting for loading dataset...")
#     with open('./found_data_set.pickle','rb') as file:
#         data_set = pickle.load(file)
#     print(f"Dataset loaded")
#     return data_set

COMMANDS = {
    "all": "Export full dataset",
    "attitude": "Export attitude control dataset"
}

if __name__ == "__main__":

    if len(sys.argv) == 1:
        my_tools.print_help(sys.argv[0], COMMANDS)

    if len(sys.argv) == 2:
        if sys.argv[1] == 'all':
            dataset2csv()
        elif sys.argv[1] == 'attitude':
            dataset_attitude2csv()
        else:
            my_tools.print_help(sys.argv[0], COMMANDS)
